from typing import Any, Dict

from pydantic import Field

from core.agents.base import BaseAgent
from core.llm.factory import ModelType, create_llm_provider
from core.managers.dialogue_manager import DialogueManager
from core.utils.log import mylogger


class DialogueAgent(BaseAgent):
    model_type: ModelType = Field(
        default=ModelType.DOUBAO_FLASH, description="LLM model type"
    )
    temperature: float = Field(default=0.8, description="Model temperature")
    frequency_penalty: float = Field(default=0.1, description="Frequency penalty")
    presence_penalty: float = Field(default=0.0, description="Presence penalty")

    def model_post_init(self, __context: Any) -> None:
        """Initialize dialogue agent."""
        super().model_post_init(__context)
        self.llm = create_llm_provider(model_type=self.model_type)
        self.dialogue_manager = DialogueManager(character_name=self.name, llm=self.llm, model_type=self.model_type.value)
        mylogger.info(
            f"DialogueAgent initialized for {self.name} with {self.model_type.value}"
        )

    async def process(self, input_data: Dict[str, Any]) -> str:
        if "input" not in input_data:
            raise ValueError("Input data must contain 'input' key")

        user_input = input_data.get("input", "")

        return await self.dialogue_manager.generate_response(
            user_input=user_input,
            scenario_index=input_data.get("scenario_index", 0),
            temperature=self.temperature,
            frequency_penalty=self.frequency_penalty,
            presence_penalty=self.presence_penalty,
        )
